package com.leilei;

import cn.hutool.crypto.SmUtil;
import cn.hutool.crypto.symmetric.SymmetricCrypto;
import com.alibaba.fastjson.JSON;
import org.springframework.core.MethodParameter;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.web.bind.annotation.RestControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import javax.servlet.http.HttpServletRequest;
import java.lang.annotation.Annotation;
import java.lang.reflect.Field;
import java.util.*;

/**
 * @author: lei
 * @date: 2024-07-18 16:40
 * @desc: 响应data加密
 */
@RestControllerAdvice
public class EncryptResponseBodyAdvice implements ResponseBodyAdvice<Object> {

    private final HttpServletRequest httpServletRequest;

    public EncryptResponseBodyAdvice(HttpServletRequest httpServletRequest) {
        this.httpServletRequest = httpServletRequest;
    }

    /**
     * 有请求解密注解则开启响应data加密
     *
     * @param returnType
     * @param converterType
     * @return boolean
     * @author lei
     * @date 2024-07-18 16:36:52
     */
    @Override
    public boolean supports(MethodParameter returnType, Class<? extends HttpMessageConverter<?>> converterType) {
        return (returnType.hasParameterAnnotation(DecryptBody.class)) || hasDecryptBodyAnnotation(returnType);
    }

    private boolean hasDecryptBodyAnnotation(MethodParameter returnType) {
        Annotation[][] parameterAnnotations = Objects.requireNonNull(returnType.getMethod()).getParameterAnnotations();
        for (Annotation[] annotations : parameterAnnotations) {
            for (Annotation annotation : annotations) {
                if (annotation instanceof DecryptBody) {
                    return true;
                }
            }
        }
        return false;
    }


    @Override
    public Object beforeBodyWrite(Object body, MethodParameter returnType,
                                  org.springframework.http.MediaType selectedContentType,
                                  Class<? extends HttpMessageConverter<?>> selectedConverterType,
                                  org.springframework.http.server.ServerHttpRequest request,
                                  org.springframework.http.server.ServerHttpResponse response) {
        if (body == null) {
            return null;
        }
        // 符合响应加密，必须是模型为 com.leilei.ResultVO 或者 有code data message三字段并存
        if (body.getClass().isAssignableFrom(ResultVO.class)) {
            ResultVO result = (ResultVO) body;
            if (result.getData() == null) {
                return body;
            }
            Object data = result.getData();
            String encryptBase64 = encryptData(data);
            result.setData(encryptBase64);
            return result;
        }
        if (hasRequiredFields(body.getClass())) {
            encryptDataField(body);
        }
        return body;
    }

    /**
     * 加密数据
     *
     * @param data
     * @return String
     * @author lei
     * @date 2024-07-30 09:52:41
     */
    private String encryptData(Object data) {
        // 从请求中获取加密密钥
        String secret = (String) httpServletRequest.getAttribute(DecryptBodyArgumentResolver.ENCRYPT_SECRET);
        if (secret == null) {
            throw new RuntimeException("未获取的加密密钥,加密失败");
        }
        SymmetricCrypto sm4 = SmUtil.sm4(secret.getBytes());
        return sm4.encryptBase64(JSON.toJSONString(data));
    }

    private static final List<String> REQUIRED_FIELDS = Arrays.asList("code", "data", "message");
    private static final Map<Class<?>, Boolean> FIELD_CHECK_CACHE = new HashMap<>();


    /**
     * 添加一个缓存，避免每次都判断是有对应字段
     * @author lei
     * @date 2024-07-30 09:58:35
     * @param clazz
     * @return boolean
     */
    private boolean hasRequiredFields(Class<?> clazz) {
        // 使用缓存进行检查，若不存在则进行反射检查
        return FIELD_CHECK_CACHE.computeIfAbsent(clazz, cls -> {
            // 获取类中所有的字段
            Field[] fields = cls.getDeclaredFields();
            // 使用 Stream API 检查所有必需的字段是否存在
            return REQUIRED_FIELDS.stream().allMatch(fieldName ->
                    Arrays.stream(fields).anyMatch(field -> field.getName().equals(fieldName))
            );
        });
    }

    /**
     * 加密数据
     * @author lei
     * @date 2024-07-30 09:59:10
     * @param obj
     * @return void
     */
    private void encryptDataField(Object obj) {
        Class<?> clazz = obj.getClass();
        try {
            // 获取data字段并加密
            Field dataField = clazz.getDeclaredField("data");
            dataField.setAccessible(true);
            Object data = dataField.get(obj);

            if (data != null) {
                dataField.set(obj, encryptData(data));
            }
        } catch (NoSuchFieldException | IllegalAccessException e) {
            e.printStackTrace();
        }
    }

}
